.globl classify

.text
# =====================================
# COMMAND LINE ARGUMENTS
# =====================================
# Args:
#   a0 (int)        argc
#   a1 (char**)     argv
#   a1[1] (char*)   pointer to the filepath string of m0
#   a1[2] (char*)   pointer to the filepath string of m1
#   a1[3] (char*)   pointer to the filepath string of input matrix
#   a1[4] (char*)   pointer to the filepath string of output file
#   a2 (int)        silent mode, if this is 1, you should not print
#                   anything. Otherwise, you should print the
#                   classification and a newline.
# Returns:
#   a0 (int)        Classification
# Exceptions:
#   - If there are an incorrect number of command line args,
#     this function terminates the program with exit code 31
#   - If malloc fails, this function terminates the program with exit code 26
#
# Usage:
#   main.s <M0_PATH> <M1_PATH> <INPUT_PATH> <OUTPUT_PATH>
#   
#   s1: pointer to input rows/input rows
#   s2: pointer to input cols/input cols
#   s3: pointer to input/pointer to h/pointer to o
#   s4: pointer to m0
#   s5: m0 rows
#   s6: m0 cols
#   a0: pointer to h
#   s7: pointer to m1
#   s8: m1 rows
#   s9: m1 cols
#   s10: temporary
classify:
    li t1, 5
    bne a0, t1, invalid_argnum
    
    addi sp, sp, -48
    sw ra, 0(sp)
    sw s1, 4(sp)
    sw s2, 8(sp)
    sw s3, 12(sp)
    sw s4, 16(sp)
    sw s5, 20(sp)
    sw s6, 24(sp)
    sw s7, 28(sp)
    sw s8, 32(sp)
    sw s9, 36(sp)
    sw s10, 40(sp)
    sw s0, 44(sp)
    
    mv s0, a1
	# Read pretrained m0    
    addi sp, sp, -8
    sw a1, 0(sp)
    sw a2, 4(sp)
    
    lw s3, 4(a1) # pointer to the filepath string of m0
    li a0, 4
    jal ra, malloc # row
    beq a0, x0, malloc_error
    mv s1, a0 # pointer to m0 rows
    li a0, 4
    jal ra, malloc
    beq a0, x0, malloc_error
    mv s2, a0 # pointer to m0 cols
    mv a0, s3
    mv a1, s1
    mv a2, s2
    addi sp, sp, -12
    sw a0, 0(sp)
    sw a1, 4(sp)
    sw a2, 8(sp)
    jal ra, read_matrix
    mv s4, a0
    
    lw a0, 0(sp)
    lw a1, 4(sp)
    lw a2, 8(sp)
    addi sp, sp, 12 
    
    lw s5, 0(a1)
    lw s6, 0(a2)

    # free m0: s1 s2
    mv a0, s1
    jal ra, free
    mv a0, s2
    jal ra, free
    
    lw a1, 0(sp)
    lw a2, 4(sp)
    addi sp, sp, 8
    
	# Read pretrained m1
    addi sp, sp, -8
    sw a1, 0(sp)
    sw a2, 4(sp)
    
    lw s3, 8(a1)
    li a0, 4
    jal ra, malloc
    beq a0, x0, malloc_error
    mv s1, a0 # rows
    li a0, 4
    jal ra, malloc
    beq a0, x0, malloc_error
    mv s2, a0 # cols
    mv a0, s3
    mv a1, s1
    mv a2, s2
    addi sp, sp, -12
    sw a0, 0(sp)
    sw a1, 4(sp)
    sw a2, 8(sp)
    jal ra, read_matrix
    mv s7, a0
    lw a0, 0(sp)
    lw a1, 4(sp)
    lw a2, 8(sp)
    addi sp, sp, 12 

    
    lw s8, 0(a1)
    lw s9, 0(a2)
    
    # free m1: s1 s2
    mv a0, s1
    jal ra, free
    mv a0, s2
    jal ra, free
    
    lw a1, 0(sp)
    lw a2, 4(sp)
    addi sp, sp, 8

	# Read input matrix
    addi sp, sp, -8
    sw a1, 0(sp)
    sw a2, 4(sp)
    
    lw s3, 12(a1)
    li a0, 4
    jal ra, malloc
    beq a0, x0, malloc_error
    mv s1, a0 # rows
    li a0, 4
    jal ra, malloc
    beq a0, x0, malloc_error
    mv s2, a0 # cols
    mv a0, s3
    mv a1, s1
    mv a2, s2
    jal ra, read_matrix
    mv s3, a0
    lw a1, 0(sp)
    lw a2, 4(sp)
    addi sp, sp, 8

	# Compute h = matmul(m0, input)
    addi sp, sp, -8
    sw a1, 0(sp)
    sw a2, 4(sp)
    
    # free
    mv s10, s1
    lw s1, 0(s10) # input rows
    mv a0, s10
    jal ra, free
    mv s10, s2
    lw s2, 0(s10) # input cols
    mv a0, s10
    jal ra, free
    
    mul a0, s5, s2
    slli a0, a0, 2
    jal ra, malloc # malloc h
    beq a0, x0, malloc_error
    mv a6, a0
    mv a0, s4
    mv a1, s5
    mv a2, s6
    mv a3, s3
    mv a4, s1
    mv a5, s2
    addi sp, sp, -4
    sw a6, 0(sp)
    jal ra, matmul
    lw a6, 0(sp)
    addi sp, sp, 4
    mv s10, a6
    
    lw a1, 0(sp)
    lw a2, 4(sp)
    addi sp, sp, 8
    
    # free input m0
    mv a0, s3
    jal ra, free

	# Compute h = relu(h)
    addi sp, sp, -8
    sw a1, 0(sp)
    sw a2, 4(sp)
    mv a0, s10
    mul a1, s5, s2
    jal ra, relu
    mv s3, a0 # pointer to h
    lw a1, 0(sp)
    lw a2, 4(sp)
    addi sp, sp, 8
    
	# Compute o = matmul(m1, h)
    addi sp, sp, -8
    sw a1, 0(sp)
    sw a2, 4(sp)
    mul a0, s8, s2
    slli a0, a0, 2
    jal ra, malloc
    beq a0, x0, malloc_error
    mv a6, a0
    mv a0, s7
    mv a1, s8
    mv a2, s9
    mv a3, s3
    mv a4, s5
    mv a5, s2
    addi sp, sp, -4
    sw a6, 0(sp)
    jal ra, matmul
    lw a6, 0(sp)
    addi sp, sp, 4
    mv s10, s3
    mv s3, a6 # pointer to o
    
    # free h, m1
    mv a0, s10
    jal ra, free
    
    lw a1, 0(sp)
    lw a2, 4(sp)
    addi sp, sp, 8
    
	# Write output matrix o
    addi sp, sp, -8
    sw a1, 0(sp)
    sw a2, 4(sp)
    lw a0, 16(s0)
    mv a1, s3
    mv a2, s8
    mv a3, s2
    jal ra, write_matrix
    lw a1, 0(sp)
    lw a2, 4(sp)
    addi sp, sp, 8
    
	# Compute and return argmax(o)
    addi sp, sp, -8
    sw a1, 0(sp)
    sw a2, 4(sp)
    mv a0, s3
    mul a1, s8, s2
    jal ra, argmax
    mv s10, a0
    lw a1, 0(sp)
    lw a2, 4(sp)
    addi sp, sp, 8
    bne a2, x0, no_print
    
	# If enabled, print argmax(o) and newline
    addi sp, sp, -4
    sw a0, 0(sp)
    jal ra, print_int
    li a0, '\n'
    jal ra, print_char
    lw a0, 0(sp)
    addi sp, sp, 4
    
no_print:
    mv a0, s3
    jal ra, free
    mv a0, s4
    jal ra, free
    mv a0, s7
    jal ra, free
    mv a0, s10
    
    lw ra, 0(sp)
    lw s1, 4(sp)
    lw s2, 8(sp)
    lw s3, 12(sp)
    lw s4, 16(sp)
    lw s5, 20(sp)
    lw s6, 24(sp)
    lw s7, 28(sp)
    lw s8, 32(sp)
    lw s9, 36(sp)
    lw s10, 40(sp)
    lw s0, 44(sp)
    addi sp, sp, 48
	ret
malloc_error:
    li a0, 26
    j exit
invalid_argnum:
    li a0, 31
    j exit
